import os
import shutil
from sklearn.model_selection import train_test_split


def split_dataset(input_dir, output_dir, train_ratio=0.7, val_ratio=0.2, test_ratio=0.1, random_state=42):
    """
    将图片分类数据集拆分为训练集、验证集和测试集

    参数:
        input_dir: 输入目录，包含按类别分组的图片
        output_dir: 输出目录，将在此目录下创建train/val/test子目录
        train_ratio: 训练集比例
        val_ratio: 验证集比例
        test_ratio: 测试集比例
        random_state: 随机种子
    """
    # 检查比例是否合理
    assert abs(train_ratio + val_ratio + test_ratio - 1.0) < 1e-6, "比例之和必须等于1"

    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)
    train_dir = os.path.join(output_dir, 'train')
    val_dir = os.path.join(output_dir, 'val')
    test_dir = os.path.join(output_dir, 'test')

    for dir_path in [train_dir, val_dir, test_dir]:
        os.makedirs(dir_path, exist_ok=True)

    # 遍历每个类别目录
    for class_name in os.listdir(input_dir):
        class_dir = os.path.join(input_dir, class_name)
        if not os.path.isdir(class_dir):
            continue

        # 获取该类别的所有图片文件
        all_files = [f for f in os.listdir(class_dir) if os.path.isfile(os.path.join(class_dir, f))]

        # 先拆分为训练集和临时集(验证+测试)
        train_files, temp_files = train_test_split(
            all_files,
            test_size=(val_ratio + test_ratio),
            random_state=random_state
        )

        # 再从临时集中拆分为验证集和测试集
        val_files, test_files = train_test_split(
            temp_files,
            test_size=test_ratio / (val_ratio + test_ratio),
            random_state=random_state
        )

        # 创建各类别的子目录
        for dir_path in [train_dir, val_dir, test_dir]:
            os.makedirs(os.path.join(dir_path, class_name), exist_ok=True)

        # 复制文件到相应目录
        def copy_files(files, dest_dir):
            for f in files:
                src = os.path.join(class_dir, f)
                dst = os.path.join(dest_dir, class_name, f)
                shutil.copy2(src, dst)

        copy_files(train_files, train_dir)
        copy_files(val_files, val_dir)
        copy_files(test_files, test_dir)

        print(
            f"类别 {class_name} 处理完成: 训练集 {len(train_files)}, 验证集 {len(val_files)}, 测试集 {len(test_files)}")


if __name__ == "__main__":
    # 使用示例
    input_directory = "mydata/pintu_8_new"  # 替换为你的输入目录
    output_directory = "mydata/output/pintu_8_new"  # 替换为你的输出目录

    split_dataset(input_directory, output_directory)